import os
import h5py
import torch
from torch.utils.data import Dataset, DataLoader

class AVEDataset(Dataset):
    """
    扩展的视听（AV）数据集类，支持带噪声样本的有监督学习任务
    继承自PyTorch的Dataset，用于加载干净样本、噪声样本及其对应标签
    """
    def __init__(self, data_root, split='train'):
        """
        初始化数据集
        
        参数:
            data_root (str): 数据根目录路径，包含所有特征文件和标签文件
            split (str): 数据集划分，可选'train'（训练集）、'val'（验证集）或'test'（测试集），默认为'train'
        """
        # 调用父类Dataset的初始化方法
        super(AVEDataset, self).__init__() 
        # 保存数据集划分类型
        self.split = split  
        
        # 构建各类特征和标签文件的路径
        self.visual_feature_path = os.path.join(data_root, 'visual_feature.h5')  # 视觉特征h5文件路径
        self.audio_feature_path = os.path.join(data_root, 'audio_feature.h5')    # 音频特征h5文件路径
        self.labels_path = os.path.join(data_root, 'labels.h5')                  # 标签h5文件路径
        self.sample_order_path = os.path.join(data_root, f'{split}_order.h5')    # 当前划分的样本顺序h5文件路径
         
        # 标记h5文件是否已打开，初始为False
        self.h5_isOpen = False 

    def __getitem__(self, index):
        """
        根据索引获取单个样本数据
        
        参数:
            index (int): 样本索引（基于当前划分的顺序）
            
        返回:
            visual_feat: 视觉特征
            audio_feat: 音频特征
            label: 对应的标签
        """
        # 延迟打开h5文件：首次获取样本时才打开，避免初始化时就打开所有文件
        if not self.h5_isOpen:
            # 打开h5文件并获取对应数据集
            self.visual_feature = h5py.File(self.visual_feature_path, 'r')['avadataset']
            self.audio_feature = h5py.File(self.audio_feature_path, 'r')['avadataset']
            self.labels = h5py.File(self.labels_path, 'r')['avadataset']
            self.sample_order = h5py.File(self.sample_order_path, 'r')['order']
            self.h5_isOpen = True  # 标记文件已打开
        
        # 根据当前划分的顺序文件获取实际样本索引（解决不同划分使用不同样本子集的问题）
        sample_index = self.sample_order[index]
        
        # 根据实际索引提取对应特征和标签
        visual_feat = self.visual_feature[sample_index]
        audio_feat = self.audio_feature[sample_index]
        label = self.labels[sample_index]

        return visual_feat, audio_feat, label


    def __len__(self):
        """
        获取当前划分的样本总数
        
        返回:
            sample_num (int): 样本数量
        """
        # 打开样本顺序文件，读取样本数量后立即关闭
        f = h5py.File(self.sample_order_path, 'r')
        # 'order'数据集的长度即为样本数量
        f.close()
        sample_num = len(f['order'])  
        return sample_num

